{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88f33adc-30d2-42f7-af5c-b41fbe1f6d98",
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "from moleculekit.molecule import Molecule\n",
    "from moleculekit.periodictable import periodictable\n",
    "import numpy as np\n",
    "import torch as pt\n",
    "from torch.utils.benchmark import Timer\n",
    "from torchmdnet.models.model import create_model\n",
    "\n",
    "\n",
    "# TensorNet\n",
    "model_1 = create_model({\n",
    "    'embedding_dimension': 32,\n",
    "    'num_layers': 2,\n",
    "    'num_linears_tensor': 2,\n",
    "    'num_linears_scalar': 2,\n",
    "    'num_rbf': 32,\n",
    "    'rbf_type': 'expnorm',\n",
    "    'trainable_rbf': False,\n",
    "    'activation': 'silu',\n",
    "    'neighbor_embedding': True,\n",
    "    'cutoff_lower': 0.0,\n",
    "    'cutoff_upper': 4.5,\n",
    "    'max_z': 100,\n",
    "    'max_num_neighbors': 128,\n",
    "    'model': 'tensornet',\n",
    "    'aggr': 'add',\n",
    "    'derivative': False,\n",
    "    'atom_filter': -1,\n",
    "    'prior_model': None,\n",
    "    'output_model': 'Scalar',\n",
    "    'reduce_op': 'add'\n",
    "})\n",
    "\n",
    "# ET\n",
    "model_2 = create_model({\n",
    "    'embedding_dimension': 64,\n",
    "    'attn_activation': 'silu',\n",
    "    'num_layers': 4,\n",
    "    'num_heads': 8,\n",
    "    'distance_influence': 'both',\n",
    "    'num_rbf': 32,\n",
    "    'rbf_type': 'expnorm',\n",
    "    'trainable_rbf': False,\n",
    "    'activation': 'silu',\n",
    "    'neighbor_embedding': True,\n",
    "    'cutoff_lower': 0.0,\n",
    "    'cutoff_upper': 5.0,\n",
    "    'max_z': 100,\n",
    "    'max_num_neighbors': 128,\n",
    "    'model': 'equivariant-transformer',\n",
    "    'aggr': 'add',\n",
    "    'derivative': False,\n",
    "    'atom_filter': -1,\n",
    "    'prior_model': None,\n",
    "    'output_model': 'Scalar',\n",
    "    'reduce_op': 'add'\n",
    "})\n",
    "\n",
    "\n",
    "def benchmark(model, pdb_file, device, compute_forces=True, batch_size=1):\n",
    "\n",
    "    model = deepcopy(model).to(device)\n",
    "\n",
    "    # Get molecular data\n",
    "    molecule = Molecule(pdb_file)\n",
    "    atomic_numbers = pt.tensor([periodictable[symbol].number for symbol in molecule.element], dtype=pt.long, device=device)\n",
    "    positions = pt.tensor(molecule.coords[:,:,0], dtype=pt.float32, device=device).to(device)\n",
    "\n",
    "    # Setup a batch\n",
    "    batch = pt.flatten(pt.tile(pt.arange(batch_size).unsqueeze(1), (1, len(atomic_numbers)))).to(device)\n",
    "    atomic_numbers = pt.tile(atomic_numbers, (batch_size,))\n",
    "    positions = pt.tile(positions, (batch_size, 1)).detach()\n",
    "\n",
    "    # Setup the force computation\n",
    "    assert not (compute_forces and (batch_size > 1))\n",
    "    positions.requires_grad = compute_forces\n",
    "    \n",
    "    # Benchmark\n",
    "    stmt = f'''\n",
    "        energy = model(atomic_numbers, positions, batch)\n",
    "        {'energy[0].sum().backward()' if compute_forces else ''}\n",
    "        '''\n",
    "    timer = Timer(stmt=stmt, globals=locals())\n",
    "    speed = timer.blocked_autorange(min_run_time=10).mean * 1000 # s --> ms\n",
    "\n",
    "    return speed\n",
    "\n",
    "# Benchmarking speed\n",
    "device = pt.device('cuda:0')\n",
    "systems = [('/systems/alanine_dipeptide.pdb', 'ALA2'),\n",
    "           ('/systems/chignolin.pdb', 'CLN'),\n",
    "           ('/systems/dhfr.pdb', 'DHFR'),\n",
    "           ('/systems/factorIX.pdb', 'FC9')]\n",
    "\n",
    "methods = [('TensorNet', model_1), ('ET', model_2)]\n",
    "\n",
    "speed_methods = {}\n",
    "for meth, model in methods:\n",
    "    speed_methods[meth] = {}\n",
    "    print(f'Method: {meth}')\n",
    "    for pdb_file, name in systems:\n",
    "        try:\n",
    "            speed = benchmark(model, pdb_file, device, compute_forces=True, batch_size=1)\n",
    "            speed_methods[meth][name] = speed\n",
    "            print(f'  {name}: {speed} ms/it')\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            print(f'  {name}: failed')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
